from torch.utils.data import Dataset
import pandas as pd
import numpy as np
import torch
import os

class iris_dataloader(Dataset):
    def __init__(self, data_path):
        # 检查路径是否存在
        self.data_path = data_path
        assert os.path.exists(self.data_path), "data path does not exist"
        
        # 读取文件数据
        df = pd.read_csv(self.data_path, names=[0, 1, 2, 3, 4])
        d = {"setosa":0, "versicolor":1, "virginica":2}
        df[4] = df[4].map(d)
        self.data = df.iloc[1:, :4]
        self.label = df.iloc[1:, 4]
        
        # 转换为numpy数组
        self.data = np.array(self.data, dtype=np.float32)
        self.label = np.array(self.label, dtype=np.int64)
        
        # 标准化处理
        self.data = (self.data - self.data.mean()) / self.data.std()
        
        # 转换为tensor
        self.data = torch.tensor(np.array(self.data))
        self.label = torch.tensor(np.array(self.label))
         
        print(f'🚀数据集加载成功!\n数据集大小为: {len(self.label)}')

    def __getitem__(self, index):
        return self.data[index], self.label[index]
    
    def __len__(self):
        return len(self.label)
    
if __name__ == '__main__':
    iris_dataloader('iris.csv')